#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DilithiumAlgorithm的泄漏模型
定义Dilithium各个操作的Power泄漏特征
"""

import numpy as np
from typing import Dict, Any, List, Tuple


class DilithiumLeakageModels:
    """Dilithium泄漏模型集合"""
    
    # DilithiumParameters（Dilithium2）
    DILITHIUM_Q = 8380417  # 模数
    DILITHIUM_D = 13       # 舍入Parameters
    DILITHIUM_TAU = 39     # 挑战权重
    DILITHIUM_GAMMA1 = (1 << 17)  # 系数范围
    DILITHIUM_GAMMA2 = ((DILITHIUM_Q - 1) // 88)  # Hint范围
    DILITHIUM_K = 4        # Matrix维度
    DILITHIUM_L = 4        # Matrix维度
    DILITHIUM_ETA = 2      # Key系数范围
    
    @staticmethod
    def hamming_weight(value: int) -> int:
        """Calculate汉明权重"""
        return bin(value).count('1')
    
    @staticmethod
    def hamming_distance(a: int, b: int) -> int:
        """CalculateHamming distance"""
        return DilithiumLeakageModels.hamming_weight(a ^ b)
    
    @classmethod
    def ntt_transform_model(cls, coeffs: List[int], q: int = None, omega: int = None) -> List[int]:
        """
        数论变换(NTT)泄漏模型（带旋转因子）
        
        Args:
            coeffs: 多项式系数
            q: 模数（DefaultUseDilithium的Q）
            omega: 原根（DefaultUseDilithium的原根）
            
        Returns:
            NTT变换后的系数
            
        Description:
        NTTYesFFTInHas限域上的实现，Require旋转因子（twiddle factors）
        ForDilithium，Useq=8380417，omegaFor256阶原根
        
        Note：这Yes教学用简化版，实际DilithiumUse优化的NTT实现
        """
        if q is None:
            q = cls.DILITHIUM_Q
        
        n = len(coeffs)
        if n <= 1:
            return coeffs
        
        # Dilithium的256阶原根（简化：这里Use理论值）
        # 实际DilithiumUse预Calculate的旋转因子表
        if omega is None:
            omega = 1753  # Dilithium的256阶原根之一
        
        # Calculate旋转因子表（简化版）
        def compute_twiddle_factors(n: int, omega: int, q: int) -> List[int]:
            """Calculate旋转因子"""
            factors = [1]
            for i in range(1, n):
                factors.append((factors[-1] * omega) % q)
            return factors
        
        twiddle_factors = compute_twiddle_factors(n, omega, q)
        
        # 简化的Cooley-Tukey NTT（带旋转因子）
        result = coeffs.copy()
        
        # 蝶形运算
        length = 2
        twiddle_idx_step = n // 2
        
        while length <= n:
            half_len = length // 2
            twiddle_step = n // length
            
            for i in range(0, n, length):
                twiddle_idx = 0
                for j in range(half_len):
                    # 获取旋转因子
                    w = twiddle_factors[twiddle_idx]
                    
                    u = result[i + j]
                    v = (result[i + j + half_len] * w) % q
                    
                    result[i + j] = (u + v) % q
                    result[i + j + half_len] = (u - v + q) % q
                    
                    twiddle_idx = (twiddle_idx + twiddle_step) % n
            
            length *= 2
        
        return result
    
    @classmethod
    def inv_ntt_transform_model(cls, coeffs: List[int], q: int = None) -> List[int]:
        """
        逆NTT变换泄漏模型
        
        Args:
            coeffs: NTT域的系数
            q: 模数
            
        Returns:
            时域系数
        """
        if q is None:
            q = cls.DILITHIUM_Q
        
        n = len(coeffs)
        
        # UseNTT的共轭性质Calculate逆变换
        # inv_NTT(x) = (1/n) * NTT(x, omega^(-1))
        
        # 简化：直接Return原系数（实际应该Complete实现）
        # TODO: 实现Complete的逆NTT
        return coeffs
    
    @classmethod
    def polynomial_multiplication_model(cls, a: List[int], b: List[int], q: int) -> List[int]:
        """
        多项式乘法泄漏模型
        
        Args:
            a, b: 多项式系数
            q: 模数
            
        Returns:
            乘积多项式系数
        """
        n = max(len(a), len(b))
        # 简化的多项式乘法
        result = [0] * (2 * n - 1)
        
        for i in range(len(a)):
            for j in range(len(b)):
                result[i + j] = (result[i + j] + a[i] * b[j]) % q
        
        return result[:n]  # 截断To原长度
    
    @classmethod
    def rejection_sampling_model(cls, z: List[int], bound: int) -> bool:
        """
        拒绝Sampling泄漏模型
        
        Args:
            z: 候选Vector
            bound: 边界值
            
        Returns:
            YesNo接受
        """
        # CalculateNo穷范数
        max_coeff = max(abs(coeff) for coeff in z)
        return max_coeff < bound
    
    @classmethod
    def hint_generation_model(cls, w: List[int], r: List[int], q: int, gamma2: int) -> List[int]:
        """
        HintGenerate泄漏模型
        
        Args:
            w: wVector
            r: rVector
            q: 模数
            gamma2: Parameters
            
        Returns:
            HintVector
        """
        hints = []
        for i in range(len(w)):
            # 简化的HintGenerate
            if abs(r[i]) > gamma2:
                hints.append(1)
            else:
                hints.append(0)
        
        return hints
    
    @classmethod
    def challenge_generation_model(cls, commitment: bytes, message: bytes) -> List[int]:
        """
        挑战Generate泄漏模型
        
        Args:
            commitment: 承诺值
            message: 消息
            
        Returns:
            挑战多项式
        """
        # 简化的挑战Generate（实际UseSHAKE256）
        seed = commitment + message
        challenge = []
        
        # 简化的伪随机Generate
        for i in range(256):  # Dilithium多项式长度
            val = (sum(seed) + i) % 3 - 1  # Generate{-1, 0, 1}
            challenge.append(val)
        
        return challenge
    
    @classmethod
    def power_of_two_rounding_model(cls, r: int, d: int, q: int) -> Tuple[int, int]:
        """
        2的幂次舍入泄漏模型
        
        Args:
            r: Input值
            d: 舍入Parameters
            q: 模数
            
        Returns:
            (r1, r0) 高位And低位
        """
        r = r % q
        r0 = r % (1 << d)
        if r0 > (1 << (d - 1)):
            r0 -= (1 << d)
        
        r1 = (r - r0) // (1 << d)
        return (r1, r0)
    
    @classmethod
    def decompose_model(cls, r: int, q: int, gamma2: int) -> Tuple[int, int]:
        """
        分解泄漏模型
        
        Args:
            r: Input值
            q: 模数
            gamma2: Parameters
            
        Returns:
            (r1, r0) 分解Result
        """
        r = r % q
        r0 = r % (2 * gamma2)
        if r0 > gamma2:
            r0 -= 2 * gamma2
        
        r1 = (r - r0) // (2 * gamma2)
        return (r1, r0)
    
    @classmethod
    def signature_generation_model(cls, message: bytes, private_key: Dict[str, Any],
                                 nonce: bytes) -> Tuple[List[int], List[int], List[int]]:
        """
        Dilithium签名Generate泄漏模型（简化版）
        
        Args:
            message: 消息
            private_key: 私钥
            nonce: 随机数
            
        Returns:
            签名(c, z, h)
        """
        # 简化的签名Generate过程
        # 实际Dilithium签名更复杂
        
        # 1. Generatey
        y = [int.from_bytes(nonce[i:i+4], 'little') % cls.DILITHIUM_Q 
             for i in range(0, min(len(nonce), 1024), 4)]
        if len(y) < 256:
            y.extend([0] * (256 - len(y)))
        y = y[:256]
        
        # 2. Calculatew = Ay（简化）
        w = [(sum(y) + i) % cls.DILITHIUM_Q for i in range(256)]
        
        # 3. Generate挑战c
        commitment = bytes(w[:32])  # 简化的承诺
        c = cls.challenge_generation_model(commitment, message)
        
        # 4. Calculatez = y + cs（简化）
        s = private_key.get('s', [1] * 256)  # 简化的私钥
        z = [(y[i] + c[i] * s[i]) % cls.DILITHIUM_Q for i in range(256)]
        
        # 5. GenerateHinth（简化）
        h = cls.hint_generation_model(w, z, cls.DILITHIUM_Q, cls.DILITHIUM_GAMMA2)
        
        return (c, z, h)
    
    @classmethod
    def verification_model(cls, message: bytes, signature: Tuple[List[int], List[int], List[int]],
                          public_key: Dict[str, Any]) -> bool:
        """
        DilithiumVerify泄漏模型（简化版）
        
        Args:
            message: 消息
            signature: 签名(c, z, h)
            public_key: 公钥
            
        Returns:
            VerifyResult
        """
        c, z, h = signature
        
        # 简化的Verify过程
        # 1. Checkz的范数
        if not cls.rejection_sampling_model(z, cls.DILITHIUM_GAMMA1):
            return False
        
        # 2. 重新Calculate挑战（简化）
        w_prime = [(sum(z) - c[i]) % cls.DILITHIUM_Q for i in range(len(c))]
        commitment = bytes(w_prime[:32])
        c_prime = cls.challenge_generation_model(commitment, message)
        
        # 3. Compare挑战
        return c == c_prime
    
    @classmethod
    def get_intermediate_functions(cls) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return {
            'ntt_transform': cls.ntt_transform_model,
            'inv_ntt_transform': cls.inv_ntt_transform_model,
            'polynomial_multiplication': cls.polynomial_multiplication_model,
            'rejection_sampling': cls.rejection_sampling_model,
            'hint_generation': cls.hint_generation_model,
            'challenge_generation': cls.challenge_generation_model,
            'power_of_two_rounding': cls.power_of_two_rounding_model,
            'decompose': cls.decompose_model,
            'signature_generation': cls.signature_generation_model,
            'verification': cls.verification_model,
        }
    
    @classmethod
    def get_leakage_models(cls) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return {
            'value': lambda x: x,
            'hamming_weight': cls.hamming_weight,
            'hamming_distance': cls.hamming_distance,
            'bit_value': lambda x, bit_pos: (x >> bit_pos) & 1,
        }
    
    @classmethod
    def get_target_operations(cls) -> List[str]:
        """获取可Attack的目标操作"""
        return [
            'ntt_transform',
            'polynomial_multiplication',
            'rejection_sampling',
            'hint_generation',
            'challenge_generation',
            'signature_generation',
            'verification',
            'rounding_operation',
            'decomposition',
            'matrix_vector_mult'
        ]
    
    @classmethod
    def get_operation_to_intermediate_mapping(cls) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return {
            'ntt_transform': 'ntt_transform',
            'polynomial_multiplication': 'polynomial_multiplication',
            'rejection_sampling': 'rejection_sampling',
            'hint_generation': 'hint_generation',
            'challenge_generation': 'challenge_generation',
            'signature_generation': 'signature_generation',
            'verification': 'verification',
            'rounding_operation': 'power_of_two_rounding',
            'decomposition': 'decompose',
            'matrix_vector_mult': 'polynomial_multiplication',
        }
    
